// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/exec/aggregate.h>
#include <pollux/exec/set_accumulator.h>

namespace kumo::pollux::functions::aggregate {
    /// @tparam ignoreNulls Whether null inputs are ignored.
    /// @tparam nullForEmpty When true, nulls are returned for empty groups.
    /// Otherwise, empty arrays.
    template<
        typename T,
        bool ignoreNulls = false,
        bool nullForEmpty = true,
        typename AccumulatorType = pollux::aggregate::prestosql::SetAccumulator<T> >
    class SetBaseAggregate : public exec::Aggregate {
    public:
        explicit SetBaseAggregate(const TypePtr &resultType)
            : exec::Aggregate(resultType) {
        }

        int32_t accumulatorFixedWidthSize() const override {
            return sizeof(AccumulatorType);
        }

        bool isFixedSize() const override {
            return false;
        }

        void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
        override {
            auto array_vector = (*result)->as<ArrayVector>();
            array_vector->resize(numGroups);

            auto *rawOffsets = array_vector->offsets()->asMutable<vector_size_t>();
            auto *rawSizes = array_vector->sizes()->asMutable<vector_size_t>();

            vector_size_t numValues = 0;
            uint64_t *rawNulls = getRawNulls(array_vector);
            for (auto i = 0; i < numGroups; ++i) {
                auto *group = groups[i];
                if (isNull(group)) {
                    if constexpr (nullForEmpty) {
                        array_vector->set_null(i, true);
                    } else {
                        // If the group's accumulator is null, the corresponding result is an
                        // empty array.
                        clearNull(rawNulls, i);
                        rawOffsets[i] = 0;
                        rawSizes[i] = 0;
                    }
                } else {
                    clearNull(rawNulls, i);

                    const auto size = value(group)->size();

                    rawOffsets[i] = numValues;
                    rawSizes[i] = size;

                    numValues += size;
                }
            }

            if constexpr (std::is_same_v<T, ComplexType>) {
                auto values = array_vector->elements();
                values->resize(numValues);

                vector_size_t offset = 0;
                for (auto i = 0; i < numGroups; ++i) {
                    auto *group = groups[i];
                    if (!isNull(group)) {
                        offset += value(group)->extractValues(*values, offset);
                    }
                }
            } else {
                auto values = array_vector->elements()->as<FlatVector<T> >();
                values->resize(numValues);

                vector_size_t offset = 0;
                for (auto i = 0; i < numGroups; ++i) {
                    auto *group = groups[i];
                    if (!isNull(group)) {
                        offset += value(group)->extractValues(*values, offset);
                    }
                }
            }
        }

        void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
        override {
            return extractValues(groups, numGroups, result);
        }

        void addIntermediateResults(
            char **groups,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool /*mayPushdown*/) override {
            addIntermediateResultsInt(groups, rows, args, false);
        }

        void addSingleGroupIntermediateResults(
            char *group,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool /*mayPushdown*/) override {
            addSingleGroupIntermediateResultsInt(group, rows, args, false);
        }

    protected:
        inline AccumulatorType *value(char *group) {
            return reinterpret_cast<AccumulatorType *>(group + Aggregate::offset_);
        }

        void addIntermediateResultsInt(
            char **groups,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool clearNullForAllInputs) {
            decoded_.decode(*args[0], rows);

            auto baseArray = decoded_.base()->template as<ArrayVector>();
            decodedElements_.decode(*baseArray->elements());

            rows.applyToSelected([&](vector_size_t i) {
                if (decoded_.is_null_at(i)) {
                    if (clearNullForAllInputs) {
                        clearNull(groups[i]);
                    }
                    return;
                }

                auto *group = groups[i];
                clearNull(group);

                auto tracker = trackRowSize(group);

                auto decodedIndex = decoded_.index(i);
                if constexpr (ignoreNulls) {
                    value(group)->addNonNullValues(
                        *baseArray, decodedIndex, decodedElements_, allocator_);
                } else {
                    value(group)->addValues(
                        *baseArray, decodedIndex, decodedElements_, allocator_);
                }
            });
        }

        void addSingleGroupIntermediateResultsInt(
            char *group,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool clearNullForAllInputs) {
            decoded_.decode(*args[0], rows);

            auto baseArray = decoded_.base()->template as<ArrayVector>();

            decodedElements_.decode(*baseArray->elements());

            auto *accumulator = value(group);

            auto tracker = trackRowSize(group);
            rows.applyToSelected([&](vector_size_t i) {
                if (decoded_.is_null_at(i)) {
                    if (clearNullForAllInputs) {
                        clearNull(group);
                    }
                    return;
                }

                clearNull(group);

                auto decodedIndex = decoded_.index(i);
                if constexpr (ignoreNulls) {
                    accumulator->addNonNullValues(
                        *baseArray, decodedIndex, decodedElements_, allocator_);
                } else {
                    accumulator->addValues(
                        *baseArray, decodedIndex, decodedElements_, allocator_);
                }
            });
        }

        void initializeNewGroupsInternal(
            char **groups,
            melon::Range<const vector_size_t *> indices) override {
            const auto &type = resultType()->childAt(0);
            exec::Aggregate::setAllNulls(groups, indices);
            for (auto i: indices) {
                new(groups[i] + offset_) AccumulatorType(type, allocator_);
            }
        }

        void destroyInternal(melon::Range<char **> groups) override {
            for (auto *group: groups) {
                if (isInitialized(group) && !isNull(group)) {
                    value(group)->free(*allocator_);
                }
            }
        }

        DecodedVector decoded_;
        DecodedVector decodedElements_;
    };

    template<
        typename T,
        bool ignoreNulls = false,
        bool nullForEmpty = true,
        typename AccumulatorType = pollux::aggregate::prestosql::SetAccumulator<T> >
    class SetAggAggregate
            : public SetBaseAggregate<T, ignoreNulls, nullForEmpty, AccumulatorType> {
    public:
        using Base = SetBaseAggregate<T, ignoreNulls, nullForEmpty, AccumulatorType>;

        explicit SetAggAggregate(const TypePtr &resultType) : Base(resultType) {
        }

        bool supportsToIntermediate() const override {
            return true;
        }

        void toIntermediate(
            const SelectivityVector &rows,
            std::vector<VectorPtr> &args,
            VectorPtr &result) const override {
            const auto &elements = args[0];
            const auto numRows = rows.size();

            // Convert input to a single-entry array.

            // Set nulls for rows not present in 'rows'.
            auto *pool = Base::allocator_->pool();
            BufferPtr nulls = allocate_nulls(numRows, pool);
            memcpy(
                nulls->asMutable<uint64_t>(),
                rows.asRange().bits(),
                bits::nbytes(numRows));

            // Set offsets to 0, 1, 2, 3...
            BufferPtr offsets = allocateOffsets(numRows, pool);
            auto *rawOffsets = offsets->asMutable<vector_size_t>();
            std::iota(rawOffsets, rawOffsets + numRows, 0);

            // Set sizes to 1.
            BufferPtr sizes = allocateSizes(numRows, pool);
            auto *rawSizes = sizes->asMutable<vector_size_t>();
            std::fill(rawSizes, rawSizes + numRows, 1);

            result = std::make_shared<ArrayVector>(
                pool,
                ARRAY(elements->type()),
                nulls,
                numRows,
                offsets,
                sizes,
                BaseVector::loaded_vector_shared(elements));
        }

        void addRawInput(
            char **groups,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool /*mayPushdown*/) override {
            Base::decoded_.decode(*args[0], rows);
            rows.applyToSelected([&](vector_size_t i) {
                auto *group = groups[i];
                Base::clearNull(group);

                auto tracker = Base::trackRowSize(group);
                if constexpr (ignoreNulls) {
                    Base::value(group)->addNonNullValue(
                        Base::decoded_, i, Base::allocator_);
                } else {
                    Base::value(group)->addValue(Base::decoded_, i, Base::allocator_);
                }
            });
        }

        void addSingleGroupRawInput(
            char *group,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool /*mayPushdown*/) override {
            Base::decoded_.decode(*args[0], rows);

            Base::clearNull(group);
            auto *accumulator = Base::value(group);

            auto tracker = Base::trackRowSize(group);
            rows.applyToSelected([&](vector_size_t i) {
                if constexpr (ignoreNulls) {
                    accumulator->addNonNullValue(Base::decoded_, i, Base::allocator_);
                } else {
                    accumulator->addValue(Base::decoded_, i, Base::allocator_);
                }
            });
        }
    };
} // namespace kumo::pollux::functions::aggregate
